import pickle
import threading
import traceback
from concurrent.futures import ThreadPoolExecutor
from typing import Optional

import msgpack
import zmq

from ctaf_core.ctaf_objects.base_service_request import BaseServiceReq
from ctaf_core.ctaf_objects.base_service_response import BaseServiceResp
from ctaf_core.ctaf_utils.config_process_util import load_bus_service_sever_settings
from ctaf_facade.enumerations.service_rsp_code_enum import ResponseCodeEnum

from ctaf_core.bus_server_bundle.bus_service_container import BusServiceContainer, BusServiceContainerFactory
from ctaf_core.ctaf_objects.base_event import MsgEvent, msgpack_event_default
from ctaf_core.log_trace_bundle.logging import LogRecorder, LoggerFactory


class BusServiceServer:

    log_recorder: LogRecorder = LoggerFactory.get_logger()

    # Number of worker threads for RPC processing
    WORKER_THREADS = 5

    def __init__(self, bus_service_server_setting_filename: str):

        self.bus_service_server_setting = load_bus_service_sever_settings(bus_service_server_setting_filename)

        if not self.bus_service_server_setting:
            return

        self.bus_server_name = self.bus_service_server_setting.get("busServerName")

        # Address for clients to connect to the ROUTER
        self.bus_server_router_address = self.bus_service_server_setting.get(
            "busServerRepAddress")

        # Internal address for workers to connect to the ROUTER
        self.bus_server_worker_address = f"inproc://{self.bus_server_name}_workers"

        self.bus_server_pub_address = self.bus_service_server_setting.get("busServerPubAddress")

        self.sevice_container: BusServiceContainer = BusServiceContainerFactory.get_bus_service_container(
            self.bus_service_server_setting.get("rpcServiceInfo"))

        # Zmq port related
        self._context: zmq.Context = zmq.Context()

        # ROUTER socket for clients (frontend)
        self._socket_router: zmq.Socket = self._context.socket(zmq.ROUTER)
        BusServiceServer.sockcet_alive_set(self._socket_router)

        # DEALER socket for workers (backend)
        self._socket_dealer_backend: zmq.Socket = self._context.socket(zmq.DEALER)
        BusServiceServer.sockcet_alive_set(self._socket_dealer_backend)

        # Publish socket (Publish–subscribe pattern)
        self._socket_pub: zmq.Socket = self._context.socket(zmq.PUB)
        BusServiceServer.sockcet_alive_set(self._socket_pub)

        # Worker thread pool related
        self._active: bool = False
        self._worker_executor: Optional[ThreadPoolExecutor] = None
        self._lock: threading.Lock = threading.Lock()

        # Flag to signal workers to terminate gracefully
        self._worker_active: threading.Event = threading.Event()
        self._worker_active.set()

        #frontend-backend proxy thread
        self._proxy_thread: Optional[threading.Thread] = None

    def is_active(self) -> bool:
        """"""
        return self._active

    def start(self) -> None:
        """
        Start BusServiceServer
        """
        if self._active:
            return

        # Bind ROUTER socket for clients
        self._socket_router.bind(self.bus_server_router_address)
        BusServiceServer.log_recorder.write_info_log(
            f"BusServiceServer ROUTER bound to {self.bus_server_router_address} for clients.")

        # Bind DEALER backend socket for workers
        self._socket_dealer_backend.bind(self.bus_server_worker_address)
        BusServiceServer.log_recorder.write_info_log(
            f"BusServiceServer DEALER backend bound to {self.bus_server_worker_address} for workers.")

        # Bind PUB socket
        self._socket_pub.bind(self.bus_server_pub_address)
        BusServiceServer.log_recorder.write_info_log(
            f"BusServiceServer PUB bound to {self.bus_server_pub_address}.")

        # Set BusServiceServer status
        self._active = True

        # Start the ROUTER-DEALER proxy in a separate thread
        self._proxy_thread = threading.Thread(target=self._run_proxy)
        self._proxy_thread.start()
        BusServiceServer.log_recorder.write_info_log("BusServiceServer ROUTER-DEALER proxy started.")

        # Initialize and start ThreadPoolExecutor for workers
        self._worker_executor = ThreadPoolExecutor(max_workers=self.WORKER_THREADS,
                                                   thread_name_prefix="BusSeverRPCWorker")
        for i in range(self.WORKER_THREADS):
            self._worker_executor.submit(self._worker_task, i)
            BusServiceServer.log_recorder.write_info_log(
                f"BusServiceServer RPC worker {i} submitted to ThreadPoolExecutor.")

        BusServiceServer.log_recorder.write_info_log(
            "BusServiceServer started with ROUTER-DEALER proxy and workers' pool.")

    def _run_proxy(self) -> None:

        try:
            zmq.proxy(self._socket_router, self._socket_dealer_backend)
            BusServiceServer.log_recorder.write_info_log("BusServiceServer ROUTER-DEALER  proxy  exited.")

        except zmq.ContextTerminated:
            BusServiceServer.log_recorder.write_info_log("BusServiceServer ROUTER-DEALER proxy ZMQ Context terminated.")
        except zmq.error.ZMQError:
            BusServiceServer.log_recorder.write_info_log("BusServiceServer ROUTER-DEALER proxy ZMQ Socket terminated.")
        except Exception as e:
            BusServiceServer.log_recorder.write_error_log(
                f"An error occurred in BusServiceServer ROUTER-DEALER proxy: {e}\n{traceback.format_exc()}")
        finally:
            self._active = False

    def _worker_task(self, worker_id: int) -> None:
        """
        Worker task to handle RPC requests. Each worker runs in its own thread from the ThreadPoolExecutor.
        Workers connect to the DEALER backend socket.
        """
        context = self._context
        worker_socket = context.socket(zmq.DEALER)
        BusServiceServer.sockcet_alive_set(worker_socket)

        # Workers connect to the internal DEALER backend address
        try:
            worker_socket.connect(self.bus_server_worker_address)
            BusServiceServer.log_recorder.write_info_log(
                f"Worker {worker_id} connected to {self.bus_server_worker_address}.")
        except zmq.ZMQError as e:
            BusServiceServer.log_recorder.write_error_log(
                f"Worker {worker_id} failed to connect to {self.bus_server_worker_address}: {e}.")
            worker_socket.close()
            context.term()
            return

        while self._worker_active.is_set():
            try:
                if worker_socket.poll(100, zmq.POLLIN):  # 100ms timeout
                    BusServiceServer.log_recorder.write_info_log(
                        f"Worker {worker_id} detected message on DEALER socket. Attempting to receive parts..."
                    )

                    #接收消息: [客户端ID, 实际数据]
                    msg = worker_socket.recv_multipart()

                    #解析消息组成
                    client_id = msg[0]  # 客户端标识
                    request_data = msg[-1]  # 实际数据

                    request_obj = pickle.loads(request_data)
                    BusServiceServer.log_recorder.write_info_log(
                        f"Worker {worker_id} processing: [Client: {client_id}] {request_obj}"
                    )
                    rpc_service_name = ""
                    trace_id = "N/A"

                    if isinstance(request_obj, BaseServiceReq):
                        rpc_service_name = request_obj.get_called_service_name()
                        trace_id = request_obj.get_ctaf_trace_id()

                    rep: Optional[BaseServiceResp] = None

                    try:
                        rep = self.sevice_container.invoke_rpc_sevice(rpc_service_name, request_obj)
                        if isinstance(request_obj, BaseServiceReq):
                            rep.set_called_service_name(request_obj.get_called_service_name())
                            rep.set_ctaf_trace_id(request_obj.get_ctaf_trace_id())
                        rep.set_response_code(ResponseCodeEnum.SUCCESS_EXECUTE)
                        BusServiceServer.log_recorder.write_info_log(
                            f"Worker {worker_id} executed service [{rpc_service_name}] successfully, trace id is [{trace_id}]")
                    except Exception as e:
                        err_msg: str = traceback.format_exc()
                        BusServiceServer.log_recorder.write_error_log(
                            f"Worker {worker_id} execute service [{rpc_service_name}] failed, trace id is [{trace_id}]：{e}\n{err_msg}")
                        rep: BaseServiceResp = BaseServiceResp()
                        if isinstance(request_obj, BaseServiceReq):
                            rep.set_called_service_name(request_obj.get_called_service_name())
                            rep.set_ctaf_trace_id(request_obj.get_ctaf_trace_id())
                        rep.set_response_code(ResponseCodeEnum.FAILED_EXECUTE)

                    # 发送回代理: [客户端ID, 空帧, 响应对象]
                    worker_socket.send_multipart([
                        client_id,
                        b'',  #空帧(ROUTER-DEALER规范)
                        pickle.dumps(rep)
                    ])
                else:
                    continue  # No message in 100ms, continue polling

            except zmq.ContextTerminated:
                BusServiceServer.log_recorder.write_info_log(
                    f"Worker {worker_id} ZMQ Context terminated, exiting worker loop.")
                break
            except Exception as e:
                BusServiceServer.log_recorder.write_error_log(
                    f"An unexpected error occurred in worker {worker_id}: {e}\n{traceback.format_exc()}")

        worker_socket.close()
        BusServiceServer.log_recorder.write_info_log(f"Worker {worker_id} thread stopped.")

    def _cleanup_zmq_sockets(self):
        # Close the ROUTER and DEALER backend sockets
        if self._socket_router and not self._socket_router.closed:
            try:
                self._socket_router.unbind(self.bus_server_router_address)
                self._socket_router.close()
                BusServiceServer.log_recorder.write_info_log(
                    f"BusServiceServer ZMQ ROUTER unbound and closed from {self.bus_server_router_address}.")
            except Exception as e:
                BusServiceServer.log_recorder.write_error_log(f"Error unbinding/closing ROUTER socket: {e}")

        if self._socket_dealer_backend and not self._socket_dealer_backend.closed:
            try:
                self._socket_dealer_backend.unbind(self.bus_server_worker_address)
                self._socket_dealer_backend.close()
                BusServiceServer.log_recorder.write_info_log(
                    f"BusServiceServer ZMQ DEALER backend unbound and closed from {self.bus_server_worker_address}.")
            except Exception as e:
                BusServiceServer.log_recorder.write_error_log(f"Error unbinding/closing DEALER backend socket: {e}")

        if self._socket_pub and not self._socket_pub.closed:
            try:
                self._socket_pub.unbind(self.bus_server_pub_address)
                self._socket_pub.close()
                BusServiceServer.log_recorder.write_info_log(
                    f"BusServiceServer ZMQ Publisher unbound and closed from {self.bus_server_pub_address}.")
            except Exception as e:
                BusServiceServer.log_recorder.write_error_log(
                    f"Error unbinding/closing busServiceServer PUB socket: {e}")

        if self._context and not self._context.closed:
            self._context.term()
            self._context = None
            BusServiceServer.log_recorder.write_info_log("BusServiceServer ZMQ Context terminated.")

    def stop(self) -> None:
        """
        Stop BusServiceServer
        """
        if not self._active:
            return

        # Signal workers to stop
        self._worker_active.clear()
        BusServiceServer.log_recorder.write_info_log("Signaling busServiceServer worker threads to stop...")

        # Stop BusServiceServer status
        self._active = False

        # Shut down the ThreadPoolExecutor gracefully
        if self._worker_executor:
            BusServiceServer.log_recorder.write_info_log("Shutting down BusServiceServer worker threadPoolExecutor...")
            self._worker_executor.shutdown(wait=True)
            BusServiceServer.log_recorder.write_info_log("BusServiceServer worker threadPoolExecutor shut down.")
            self._worker_executor = None

        BusServiceServer.log_recorder.write_info_log("---------------------------------")

        # Join the proxy thread
        if self._proxy_thread and self._proxy_thread.is_alive():
            BusServiceServer.log_recorder.write_info_log("Joining ROUTER-DEALER proxy thread...")
            self._proxy_thread.join(timeout=1)  # Give it a moment to shut down
            if self._proxy_thread.is_alive():
                BusServiceServer.log_recorder.write_warn_log("ROUTER-DEALER proxy thread did not terminate gracefully.")
            self._proxy_thread = None

        # Cleanup ZMQ sockets and context
        self._cleanup_zmq_sockets()
        BusServiceServer.log_recorder.write_info_log("BusServiceServer stopped and all threads shut down.")

    def publish(self, msg_event: MsgEvent) -> None:
        """
        Publish data
        """
        with self._lock:
            try:
                msg_bytes = msgpack.dumps(msg_event, default=msgpack_event_default)
                self._socket_pub.send(msg_bytes)
            except TypeError as te:
                BusServiceServer.log_recorder.write_error_log(
                    f"BusServiceServer failed to serialize event with Msgpack: {te}. Event Type: {msg_event.get_event_type()}")
            except Exception as e:
                BusServiceServer.log_recorder.write_error_log(f"BusServiceServer failed to put event: {e}")

    @staticmethod
    def sockcet_alive_set(socket_instance: zmq.Socket):
        socket_instance.setsockopt(zmq.TCP_KEEPALIVE, 1)
        socket_instance.setsockopt(zmq.TCP_KEEPALIVE_IDLE, 60)